import networkx as nx
import numpy as np
# from connect import connect
import matplotlib.pyplot as plt

class space:
  def __init__(self):
    pass
  def initByPos(self,pos):
    self.pos  = pos
    (self.num,z) = pos.shape 
  def fromC(self,r):
    num, = r.shape
    self.pos = np.zeros((num,2),dtype=np.float)
    self.pos[:,0] = r.real
    self.pos[:,1] = r.imag
    self.num = num
  def toC(self):
    r = self.pos[:,0]+1j*self.pos[:,1]
    return space.rol(r)

  def init(self,num):
    r = self.fromR(np.arange(num),5)
    self.fromC(r)

  def up(self,eg,ig):
    num,dim = self.pos.shape
    reg = eg.oneStep(self.pos)
    rig = ig.oneStep(self.pos)
    drig = rig - self.pos
    dreg = reg - self.pos
    r =  np.concatenate((rig,reg,self.abs(drig),self.abs(dreg)),axis=1)
    self.pos = r

  def down(self,N):
    r = self.pos - np.average(self.pos,axis=0)
    u, s, vh = np.linalg.svd(r, full_matrices=False)
    cut = u[:,:N].real
    P = np.sqrt(np.sum(s*s))
    # print(P)
    self.pos = P*cut

  def plot(self,color=None):
    # num,dim = self.pos.shape   
    # if dim>2:
    #   u, s, vh = np.linalg.svd(self.pos, full_matrices=False)
    #   r = u[:,:2].real
    # else:
    r = self.pos
    plt.scatter(r[:,0],r[:,1],s=0.5,c=color)
    plt.show()

  def nplot(self,G):
    pos = {}
    for n in G.nodes():
      pos[n] = np.array([self.pos[n,0],self.pos[n,1]])
    nx.draw_networkx(G, pos, labels=None)
    plt.show() 

  def abs(self,r):
    v = np.zeros((self.num,1),dtype=np.float)
    v[:,0] = np.linalg.norm(r, axis=1)
    return v

  def fromR(self,p,N):
    index = np.argsort(p)
    sizel, = p.shape
    l = int(np.ceil(sizel / (N*N)))
    r = np.zeros(p.shape,dtype=np.complex128)
    for i in range(N):
      pl = (2*i+1)*l
      ps = i*i*l
      for j in range(pl):
        k = ps + j
        if k<sizel:
          npp = (i+float(j)/float(pl))
          na = npp*2*np.pi
          npda = npp
          r[index[k]] = npda * np.exp(1j*na)
    return r

  def oneStep(self,c):
    self.pos = c.oneStep(self.pos)
  def toReal(self):
    self.pos = self.pos.real
  def trol(r):
    l,=r.shape
    rr = np.sum(r.real*r.real)/l
    ii = np.sum(r.imag*r.imag)/l
    ri = np.sum(r.real*r.imag)/l
    return rr,ii,ri
  
  def rol(r):
    rr,ii,ri = space.trol(r)
    a = np.array([[rr,ri],[ri,ii]])
    D,L = np.linalg.eigh(a)
    st = L[0,0]-1j*L[0,1]
    vr = r * st
    vr = (vr.real/np.sqrt(D[0])+1j*vr.imag/np.sqrt(D[1]))
    return vr

  def min(self):
    return np.min(self.pos,axis=0)
  def max(self):
    return np.max(self.pos,axis=0)
  def grid(self,N):
    p = self.max()-self.min()
    g = np.sqrt(p[0]*p[1]/N)
    return g
  def fit(self):
    N = self.num * 4
    o = self.min()
    g = self.grid(N)
    r = np.round((self.pos-o)/g).astype(np.int)
    ns = space()
    c = np.zeros((self.num,1),dtype=np.int)
    for i in range(self.num):
      z = r[i]
      for j in range(i+1,self.num):
        if r[j,0] == z[0] and r[j,1]==z[1]:
          c[j] += 1
    cg = c*g
    ns.num = self.num
    ns.pos = np.concatenate(((r*g+o),cg),axis=1)
    return ns
